package com.gitee.c0dehvb.reactive.sample.r2dbc.connection;

import com.gitee.c0dehvb.reactive.sample.r2dbc.ShardingRewriteEngine;
import io.r2dbc.spi.Batch;
import io.r2dbc.spi.Result;
import org.reactivestreams.Publisher;

import java.util.Collections;

/**
 * 分片Batch封装
 * @author LiYilin
 * @since 2022/6/15 9:57 PM
 **/
public class ShardingBatch implements Batch {
    private final Batch batch;
    private final ShardingRewriteEngine shardingRewriteEngine;

    public ShardingBatch(Batch batch, ShardingRewriteEngine shardingRewriteEngine) {
        this.batch = batch;
        this.shardingRewriteEngine = shardingRewriteEngine;
    }

    @Override
    public Batch add(String sql) {
        String actualSQL = shardingRewriteEngine.rewriteSql(sql, Collections.emptyList());
        return batch.add(actualSQL);
    }

    @Override
    public Publisher<? extends Result> execute() {
        return batch.execute();
    }
}
